#include "CudaEpiModel.h"
#include "CudaParam.h"
#include <thrust/copy.h>
#include <thrust/binary_search.h>
#include <thrust/count.h>
#include <thrust/sort.h>
#include <thrust/remove.h>
#include <sys/time.h> 
#include <thrust/sequence.h>
#include <thrust/partition.h>
#include <omp.h>
#include <pthread.h>
#include <string>
#include <ctime>
#include <fstream>
using namespace std;


struct timeval start, finish;
void srecord(){
	gettimeofday(&start,NULL);
}
void frecord(string msn){
	gettimeofday(&finish,NULL);
	double t1,t2;
	t1 = (double)start.tv_sec + ((double)start.tv_usec/1000000.0);
	t2 = (double)finish.tv_sec + ((double)finish.tv_usec/1000000.0);
	cout << msn << " is: " << (t2 - t1) << " secs" << endl;
}

int counter = 0;

struct DayThread{
	int gpu_id;
	int id;
	int vec_id;
};


void printCounter(){
	cout << "counter: " << counter << endl;
	counter++;
}

void *threadDay(void* data){
	DayThread& thread_local = *((DayThread*)data);
	int gpu_id = thread_local.gpu_id;
	cudaSetDevice(gpu_id);

	for(int iter = 0; iter < num_gpu_iter + 2; iter++){

		if(thread_local.id == 0){
			if(iter < num_gpu_iter){
				uploadVec(gpu_data[gpu_id],iter,thread_local.vec_id);
				thread_local.vec_id--;
			}
			if(thread_local.vec_id < 0) thread_local.vec_id = 2;
		}

		if(thread_local.id == 1){
			if(iter-1 >= 0 && iter-1 < num_gpu_iter){
				computeDay(gpu_data[gpu_id],iter-1,thread_local.vec_id);
				thread_local.vec_id--;
			}
			if(thread_local.vec_id < 0) thread_local.vec_id = 2;
		}

		if(thread_local.id == 2){
			if(iter-2 >= 0){
				downloadVec(gpu_data[gpu_id],iter-2,thread_local.vec_id);
				thread_local.vec_id--;
			}
			if(thread_local.vec_id < 0) thread_local.vec_id = 2;
		}


		pthread_barrier_wait(&(gpu_data[gpu_id].barr));
	}

	return NULL;
}

void *singleThreadVersion(void* arg){
	GPUSimStruct& data = *((GPUSimStruct*)arg);
	cudaSetDevice(data.gpu_id);

	// copy community
	#pragma omp parallel for
	for(int i = 0; i < model->nNumCommunities; i++){
		(*(data.host_comm_vec))[i].copyFromCommunity(model->commvec[i]);
	}
	cudaMemcpy(thrust::raw_pointer_cast(data.comm_vec->data()),thrust::raw_pointer_cast(data.host_comm_vec->data()),
		model->nNumCommunities*sizeof(dCommunity),cudaMemcpyHostToDevice);

	// copy tract
	cudaMemcpy(thrust::raw_pointer_cast(data.tract_vec->data()),&(model->tractvec[0]),
			model->nNumTractsTotal*sizeof(Tract),cudaMemcpyHostToDevice);

	int load_pos = 0;
	int compute_pos = 0;
	int unload_pos = 0;
	
	for(int iter = 0; iter < num_gpu_iter+2; iter++){

		if(iter < num_gpu_iter){
			uploadVec(data,iter,load_pos);
			load_pos--;
		}

		if(iter-1 >= 0 && iter-1 < num_gpu_iter){
			computeDay(data,iter-1,compute_pos);
			compute_pos--;
		}
		
		if(iter-2 >= 0){
			downloadVec(data,iter-2,unload_pos);
			unload_pos--;
		}

		if(load_pos < 0) load_pos = 2;
		if(compute_pos < 0) compute_pos = 2;
		if(unload_pos < 0) unload_pos = 2;
	}

	//pthread_barrier_init(&(data.barr),NULL,3);

	//DayThread day_data[3];

	//pthread_t threads[3];
	//for(int i = 0; i < 3; i++){
	//	day_data[i].gpu_id = data.gpu_id;
	//	day_data[i].id = i;
	//	day_data[i].vec_id = 0;
	//	pthread_create(&threads[i],NULL,threadDay,(void*)&(day_data[i]));
	//}

	//for(int i = 0; i < 3; i++){
	//	pthread_join(threads[i],NULL);
	//}

	return NULL;

}

void cudaDay(){

	// run prs and pri first using original code

	epimodelDayPart1();

	
	pthread_t threads[NUM_GPUS];
	for(int i = 0; i < NUM_GPUS; i++){
		pthread_create(&threads[i],NULL,singleThreadVersion,(void*)&(gpu_data[i]));
	}

	for(int i = 0; i < NUM_GPUS; i++){
		pthread_join(threads[i],NULL);
	}	

	syncCommunity();
}


// synchronize community data on all gpus
void syncCommunity(){
	#pragma omp parallel for
	for(int i = 0; i < model->nNumCommunities; i++){
		Community& comm = model->commvec[i];
		for(int j = comm.nFirstPerson; j < comm.nLastPerson; j++){
			Person& p = model->pvec[j];
			if(isInfected(p) && p.iday == -1){
				comm.ninf[p.age]++;
				comm.nEverInfected[p.age]++;
			}
		}
	}
}

void clean(){
	for(int id = 0; id < NUM_GPUS; id++){
		cudaSetDevice(id);
		delete gpu_data[id].rand_vec;
		
		
		for(int i = 0; i < 3; i++){
			delete gpu_data[id].ppl_vec[i];
			delete gpu_data[id].comm_worker_vec[i];
		}

		
		for(int i = 0; i < 2; i++){
			delete gpu_data[id].host_ppl_vec[i];
			delete gpu_data[id].host_comm_worker_vec[i];
		}

		delete gpu_data[id].comm_vec;
		delete gpu_data[id].host_comm_vec;

		delete gpu_data[id].comm_worker_start_vec;
		delete gpu_data[id].comm_worker_end_vec;

		delete gpu_data[id].home_ppl_index;
		delete gpu_data[id].worker_ppl_index;

		delete gpu_data[id].home_infect_result;
		delete gpu_data[id].worker_infect_result;

		delete gpu_data[id].iter_comm_start;
		delete gpu_data[id].iter_comm_end;
		
		cudaFree(dvload);
		cudaFree(dseasonality);
		cudaFree(dfVaccineEfficacyByAge);
		cudaFree(dVaccineData);
		cudaFree(dcpfc);
		cudaFree(dcpfa);
		cudaFree(dcphcc);
		cudaFree(dcphca);
		cudaFree(dwithdrawcdf32);
		cudaFree(dincubationcdf32);
		cudaFree(dcps);
		cudaFree(d_travel_length_cdf);
		cudaFree(d_travel_length_cdf);
		cudaFree(d_travel_pr);
		cudaFree(dnSchoolOpeningDays);
		cudaFree(dnVaccinePriorities2);
	}
}

void cudaInitialize(EpiModel& input_model){
	model = &input_model;

	srand(time(NULL));
    // transfer constant  information to constant memory
	for(int id = 0; id < NUM_GPUS; id++){
		cudaSetDevice(id);
		cudaMemcpyToSymbol(dvload,model->vload,VLOADNSUB*VLOADNDAY*sizeof(double)); // vload
		cudaMemcpyToSymbol(dseasonality,model->seasonality,MAXRUNLENGTH*sizeof(double)); // seasonality
		cudaMemcpyToSymbol(dfVaccineEfficacyByAge,model->fVaccineEfficacyByAge,TAG*sizeof(double)); // fVaccineEfficacyByAge
		cudaMemcpyToSymbol(dVaccineData,model->VaccineData,NUMVACCINES*sizeof(vaccinedatastruct)); // vaccineData
		cudaMemcpyToSymbol(dcpfc,cpfc,TAG*sizeof(double));
		cudaMemcpyToSymbol(dcpfa,cpfa,TAG*sizeof(double));
		cudaMemcpyToSymbol(dcphcc,cphcc,TAG*sizeof(double));
		cudaMemcpyToSymbol(dcphca,cphca,TAG*sizeof(double));
		cudaMemcpyToSymbol(dwithdrawcdf32,model->withdrawcdf32,3*WITHDRAWDAYS*sizeof(unsigned int)); // withdrawcef32
		cudaMemcpyToSymbol(dincubationcdf32,incubationcdf32,3*sizeof(unsigned int)); // incubationcdf32
		cudaMemcpyToSymbol(dcps,cps,10*sizeof(double)); // cps
		cudaMemcpyToSymbol(d_travel_length_cdf,travel_length_cdf,12*sizeof(double)); // travel_length_cdf
		cudaMemcpyToSymbol(d_travel_pr,travel_pr,TAG*sizeof(double)); // travel_pr
		cudaMemcpyToSymbol(dnSchoolOpeningDays,model->nSchoolOpeningDays,56*sizeof(int));
		cudaMemcpyToSymbol(dnVaccinePriorities2,model->nVaccinePriorities2,PRIORITY_LAST*sizeof(unsigned char));
	}

	double num_iter_approx = (double)(model->nNumPerson) * sizeof(Person) * 6 * 1.5 /  pow((double)2,30) / NUM_GPUS;

	num_gpu_iter = (int)num_iter_approx;
	if(num_gpu_iter == 0) num_gpu_iter = 1;
	num_gpu_iter *= 2;

	ofstream iter_info("sim_num_gpu_iter_info.txt",ios::out|ios::app);
	iter_info << "num of people: " << model->nNumPerson << " num_gpu_iter: " << num_gpu_iter << endl;
	iter_info.close();

	// allocate worker vec
	host_comm_worker_start_vec = new vector<int>(model->nNumCommunities);
	host_comm_worker_end_vec = new vector<int>(model->nNumCommunities);
	num_outside_worker = 0;
	for(int i = 0; i < model->nNumCommunities; i++){
		(*host_comm_worker_start_vec)[i] = num_outside_worker;
		num_outside_worker += model->commvec[i].workers.size();
		(*host_comm_worker_end_vec)[i] = num_outside_worker;
	}
	cout << "outside worker: " << num_outside_worker << endl;

	// copy community vector
	for(int id = 0; id < NUM_GPUS; id++){
		cudaSetDevice(id);
		gpu_data[id].gpu_id = id;
		gpu_data[id].comm_vec = new dCommunityVec(model->nNumCommunities);
		gpu_data[id].host_comm_vec = new hCommunityVec(model->nNumCommunities);
		gpu_data[id].comm_worker_start_vec = new thrust::device_vector<int>(model->nNumCommunities);
		gpu_data[id].comm_worker_end_vec = new thrust::device_vector<int>(model->nNumCommunities);
		*(gpu_data[id].comm_worker_start_vec) = *host_comm_worker_start_vec;
		*(gpu_data[id].comm_worker_end_vec) = *host_comm_worker_end_vec;
		// iter info allocation
		gpu_data[id].iter_comm_start = new vector<int>(num_gpu_iter);
		gpu_data[id].iter_comm_end = new vector<int>(num_gpu_iter);
	}	

	
	
	// set iteration info, each iteration deal with how many commuinities
	int num_comm_per_iter = model->nNumCommunities / (num_gpu_iter * NUM_GPUS);
	int comm_iter = 0;
	for(int iter = 0; iter < num_gpu_iter; iter++){
		for(int id = 0; id < NUM_GPUS; id++){
			(*gpu_data[id].iter_comm_start)[iter] = comm_iter;
			comm_iter += num_comm_per_iter;
			if(comm_iter < model->nNumCommunities) 
				(*gpu_data[id].iter_comm_end)[iter] = comm_iter;
			else
				(*gpu_data[id].iter_comm_end)[iter] = model->nNumCommunities;
		}
	}


	// find the appropriate gpu number
	num_gpu_ppl = 0;
	for(int iter = 0; iter < num_gpu_iter; iter++){
		for(int id = 0; id < NUM_GPUS; id++){
			int start_comm_idx = (*gpu_data[id].iter_comm_start)[iter];
			int end_comm_idx = (*gpu_data[id].iter_comm_end)[iter];
			int home_ppl = model->commvec[end_comm_idx-1].nLastPerson - model->commvec[start_comm_idx].nFirstPerson;
			if(home_ppl > num_gpu_ppl) num_gpu_ppl = home_ppl;
			int worker_ppl = (*host_comm_worker_end_vec)[end_comm_idx-1] - (*host_comm_worker_start_vec)[start_comm_idx];
			if(worker_ppl > num_gpu_ppl) num_gpu_ppl = worker_ppl;
		}
	}

	
	cout << "num_gpu_ppl: " << num_gpu_ppl * 6 << endl;

	if((double)num_gpu_ppl * 6 * sizeof(Person) > 1.5 * pow((double)2,30)){
		cout << "memory problem" << endl;
		cout << "num_gpu_ppl: " << num_gpu_ppl << endl;
		getchar();
	}

	for(int id = 0; id < NUM_GPUS; id++){
		cudaSetDevice(id);

		for(int i = 0; i < 3; i++){
			gpu_data[id].ppl_vec[i] = new dPersonVec(num_gpu_ppl);
			gpu_data[id].comm_worker_vec[i] = new dPersonVec(num_gpu_ppl);
		}

		for(int i = 0; i < 2; i++){
			gpu_data[id].host_ppl_vec[i] = new vector<Person>(num_gpu_ppl);
			gpu_data[id].host_comm_worker_vec[i] = new vector<Person>(num_gpu_ppl);
		}

		gpu_data[id].home_ppl_index = new thrust::device_vector<int>(num_gpu_ppl);
		gpu_data[id].worker_ppl_index = new thrust::device_vector<int>(num_gpu_ppl);
		gpu_data[id].home_infect_result = new thrust::device_vector<bool>(num_gpu_ppl);
		gpu_data[id].worker_infect_result = new thrust::device_vector<bool>(num_gpu_ppl);


		// copy tract vector
		gpu_data[id].tract_vec = new dTractVec(model->nNumTractsTotal);

		// generate random vector
		gpu_data[id].rand_vec = new dRandVec(num_gpu_ppl);
		hRandVec rand_temp(num_gpu_ppl);
		for(int i = 0; i < num_gpu_ppl; i++){rand_temp[i] = RandGen(rand(),rand(),rand(),rand());}
		(*gpu_data[id].rand_vec) = rand_temp;

	}
}

const double defaultvacceff[VACCEFFLENGTH+1] = {0,0.001,0.004,0.011,0.023,0.043,0.07,0.106,0.153,0.211,0.28,0.363,0.46,0.572,0.7,0.7,0.7,0.7,0.7,0.7,0.7,0.7,0.702,0.71,0.73,0.766,0.82,0.897,1};

void epimodelDayPart1(){
	#pragma omp parallel for
	for (int i = 0; i < model->nNumPerson; i++) {
		Person &p = model->pvec[i];
		p.prs = 1.0-p.fBaselineVES; // susceptibility multiplier
		if (isInfectious(p))
			p.pri = model->vload[p.nWhichVload][(int)(p.iday)] *
			model->seasonality[model->nTimer/2]; // infectiousness multiplier
		else
			p.pri = 0.0;
		if (isVaccinated(p)) {  // effects of the vaccine on transmission
			if (needsBoost(p)) {
				
				p.prs *= (1.0-model->VaccineData[whichVaccine(p)].VEs*defaultvacceff[p.vday]*model->fVaccineEfficacyByAge[p.age]);
				p.pri *= (1.0-model->VaccineData[whichVaccine(p)].VEi*defaultvacceff[p.vday]*model->fVaccineEfficacyByAge[p.age]);
			} else {
				p.prs *= (1.0-model->VaccineData[whichVaccine(p)].VEs*model->VaccineData[whichVaccine(p)].vacceff[p.vday]*model->fVaccineEfficacyByAge[p.age]);
				p.pri *= (1.0-model->VaccineData[whichVaccine(p)].VEi*model->VaccineData[whichVaccine(p)].vacceff[p.vday]*model->fVaccineEfficacyByAge[p.age]);
			}
		}
		if (isAntiviral(p)) {  // effects of the antiviral on transmission
			p.prs *= (1.0-model->AVEs);
			p.pri *= (1.0-model->AVEi);
		}
		if (isSymptomatic(p))
			p.pri *= 2.0;  // symptomatic people are 2 times as infectious
	}
}